#!/usr/bin/python3
# -*- coding: utf-8 -*-
import os
import struct
import multiprocessing

video_list_count = 8
# base_path = r'/rootfs/media/kasim/DataSet/RecordEx'
# out_path = r'/rootfs/media/kasim/DataSet/RecordEx'
base_path = r'/rootfs/media/kasim/DataSet/Record'
out_path = r'/rootfs/media/kasim/DataSet/Record'
video_list_bbox_count_file = r'video_list{}_bbox_count.txt'
video_list_bbox_file = r'video_list{}_bbox.dat'
video_list_size_file = r'video_list{}.txt'

PROCESS_COUNT = video_list_count

min_video_size = 1024*1024*7
max_video_size = 1024*1024*100

select_obj_count = 1
min_frame_count = 10

select_classes = ['Person', 'Face']

select_class_label = {
    'Person': 0,
    'Cat': 1,
    'Dog': 2,
    'BabyCar': 3,
    'Face': 4,
}
select_labels = set([select_class_label[select_class] for select_class in select_classes])

LOW_THRESHOLDS = [
    0.3,  # Person
    0.3,  # Cat
    0.3,  # Dog
    0.3,  # BabyCar
    0.4,  # Face
]

select_label_class = {}
for select_class, select_label in select_class_label.items():
    select_label_class[select_label] = select_class


def select_proc(file_queue, id):
    try:
        video_size_dict = {}
        with open(os.path.join(base_path, video_list_size_file.format(id)), 'r') as file:
            for video_info in file.readlines():
                video_info = video_info.split(',')
                video_size_dict[video_info[0].strip()] = int(video_info[1].strip())
        print('Proc:{}, Read File Size Finish!'.format(id))

        video_list = []
        with open(os.path.join(base_path, video_list_bbox_count_file.format(id)), 'r') as file:
            file_count = 0
            for video_info in file.readlines():
                video_info = video_info.split(',')
                # video_list.append([video_info[0].strip(), int(video_info[1].strip())])
                video_name = video_info[0].strip()
                video_size = video_size_dict.get(video_name, 0)
                video_list.append([video_name, int(video_info[1].strip()), video_size])
                file_count += 1
                # if file_count > 100:
                #     break
        print('Proc:{}, Read File BBox Count Finish!'.format(id))

        file_count = 0
        total_file_count = len(video_list)
        with open(os.path.join(base_path, video_list_bbox_file.format(id)), 'rb') as file:
            for video_name, BBox_count, video_size in video_list:
                if video_size < min_video_size or max_video_size < video_size:
                    bbox_info = file.read(28*BBox_count)
                    if len(bbox_info) < 28*BBox_count:
                        print('Proc:{}, File Count: {}/{}'.format(id, file_count, total_file_count))
                        file_queue.put([None, True, id])
                        return
                    continue
                low_frame_bbox_count = {}
                for select_label in select_labels:
                    low_frame_bbox_count[select_label] = {}
                for j in range(BBox_count):
                    bbox_info = file.read(28)
                    if len(bbox_info) < 28:
                        print('Proc:{}, File Count: {}/{}'.format(id, file_count, total_file_count))
                        file_queue.put([None, True, id])
                        return
                    bbox = struct.unpack('6i1f', bbox_info)
                    label = bbox[1]
                    if label not in select_labels:
                        continue
                    low_thr = LOW_THRESHOLDS[label]
                    frame_id = bbox[0]
                    score = bbox[6]
                    if score >= low_thr:
                        _frame_bbox_count = low_frame_bbox_count[label]
                        if frame_id not in _frame_bbox_count.keys():
                            _frame_bbox_count[frame_id] = 1
                        else:
                            _frame_bbox_count[frame_id] += 1

                select_obj_infos = {}
                for select_label, _frame_bbox_count in low_frame_bbox_count.items():
                    if len(_frame_bbox_count) > 0:
                        # _frame_bbox_count记录每一帧中满足低阈值要求的目标数
                        _frame_bbox_counts = _frame_bbox_count.values()
                        _max_bbox_count = max(_frame_bbox_counts)  # 同一帧里最大目标数
                        if 0 < _max_bbox_count <= select_obj_count:  # 只记录目标数小于等于select_obj_count的视频
                            max_select_frame_count = 0
                            for bbox_count in _frame_bbox_counts:
                                if 0 < bbox_count <= select_obj_count:  # 记录视频中目标数大于0且小于等于select_obj_count的帧数
                                    max_select_frame_count += 1
                            max_select_obj_count = _max_bbox_count  # 最大目标数
                            if max_select_obj_count > 0:
                                select_obj_infos[select_label_class[select_label]] = {'ObjCnt': max_select_obj_count, 'ObjFrmCnt': max_select_frame_count}

                if len(select_obj_infos) > 0:
                    file_queue.put([(video_name, select_obj_infos), False, id])

                file_count += 1
                if file_count % 1000 == 0:
                    print('Proc:{}, File Count: {}/{}'.format(id, file_count, total_file_count))
                    # return
    except Exception as e:
        print('Proc:{}, File Count: {}/{}'.format(id, file_count, total_file_count))
    print('Proc:{}, File Count: {}/{}'.format(id, file_count, total_file_count))
    file_queue.put([None, True, id])


def main():
    if not os.path.exists(out_path):
        os.makedirs(out_path)
    os.system('chmod a+wr {}'.format(out_path))

    file_queue = multiprocessing.Queue()
    select_video_workers = []
    for i in range(PROCESS_COUNT):
        select_video_workers.append(multiprocessing.Process(target=select_proc, args=(file_queue, i,)))

    for i in range(PROCESS_COUNT):
        select_video_workers[i].start()

    person_count = 0
    face_count = 0
    face_video = []
    try:
        finish_worker_count = 0
        while True:
            file_info = file_queue.get(block=True)
            if file_info is None:
                break
            video_info, finish, id = file_info
            if finish:
                print('Proc{} finish'.format(id, ))
                finish_worker_count += 1
                if PROCESS_COUNT <= finish_worker_count:
                    break
                continue
            if (video_info is None) or (len(video_info) < 1):
                continue
            video_name, select_obj_infos = video_info
            person_info = select_obj_infos.get('Person', None)
            if person_info is not None:
                if (0 < person_info['ObjCnt'] <= select_obj_count) and (person_info['ObjFrmCnt'] > 0):
                    person_count += 1
                    face_info = select_obj_infos.get('Face', None)
                    if face_info is not None:
                        if (0 < face_info['ObjCnt'] <= select_obj_count) and (face_info['ObjFrmCnt'] > 0):
                            face_count += 1
                            face_video.append(video_name)
    except Exception as e:
        print(e)

    if person_count == 0:
        noface_ratio = 0
    else:
        noface_ratio = 100 * face_count / person_count
    print('{} Person Count: {} Face Count: {}, FaceRatio: {:.03f}, NoFaceRatio: {:.03f}'.format(select_obj_count, person_count, face_count, noface_ratio, (100-noface_ratio)))

    for i in range(PROCESS_COUNT):
        select_video_workers[i].join()

    file_path = os.path.join(out_path, 'one_person_face_video_list.txt')
    with open(file_path, 'w') as file:
        face_video.sort()
        for video_name in face_video:
            file.write(video_name + '\n')
    os.system('chmod a+wr {}'.format(file_path))

    print('Select Video Workers Finish!')


if __name__ == '__main__':
    main()
